from torch import nn
import torch

"""
RNN-Cell中的参数说明：
 input_size: int, vocab_size(词元的维度大小，也可以说是词的分类大小)
 hidden_size: int,RNN隐藏状态的维度大小
 bias: bool = True,是否启用偏置
 nonlinearity: str = "tanh",RNN隐藏状态的激活函数（'relu'）
 device: Any = None,模型运行的设备
 dtype: Any = None模型输出值的类型
"""

seq_len = 4
batch_size = 1
vocab_size = 10
hidden_size = 20
num_layers = 2  # 隐藏状态层

inputs = torch.randn((batch_size, vocab_size))  # (1,10)
h0 = torch.randn((batch_size, hidden_size))  # （1,20）
# RNNCell 只获取当前的隐藏状态
# 第一层隐藏状态
rnn1 = nn.RNNCell(vocab_size, hidden_size)
fc1 = nn.Linear(hidden_size, hidden_size)
# 第二层隐藏状态
rnn2 = nn.RNNCell(hidden_size, hidden_size)
fc2 = nn.Linear(hidden_size, vocab_size)

# 第一层隐藏状态得到 隐藏状态值
h1 = rnn1(inputs, h0)
# 由 隐藏状态值 预测输出
outputs1 = fc1(h1)
# 将上一轮的输出传入 第二层隐藏状态得到 隐藏状态值
h2 = rnn2(outputs1, h0)
# 由 隐藏状态值 预测最终的输出
outputs2 = fc2(h2)

print(outputs2.shape)
